import torch.nn as nn


class ZFNet(nn.Module):
    def __init__(self):
        super(ZFNet, self).__init__()
        # n*3*227*227
        self.c1 = nn.Sequential(
            nn.Conv2d(3, 24, kernel_size=7, stride=2),
            nn.ELU(),
            nn.MaxPool2d(kernel_size=3, stride=2))
        # n*24*55*55
        self.c2 = nn.Sequential(
            nn.Conv2d(24, 64, kernel_size=5, stride=2),
            nn.ELU(),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1))
        # n*64*13*13
        self.c3 = nn.Sequential(
            nn.Conv2d(64, 96, kernel_size=3, padding=1),
            nn.ELU())
        # n*96*13*13
        self.c4 = nn.Sequential(
            nn.Conv2d(96, 96, kernel_size=3, padding=1),
            nn.ELU())
        # n*96*13*13
        self.c5 = nn.Sequential(
            nn.Conv2d(96, 64, kernel_size=3, padding=1),
            nn.ELU(),
            nn.MaxPool2d(kernel_size=3, stride=2))
        # n*64*6*6
        self.fc6 = nn.Sequential(
            nn.Linear(64 * 6 * 6, 1024),
            nn.Softsign())
        # n*1024
        self.fc7 = nn.Sequential(
            nn.Linear(1024, 1024),
            nn.Softsign())
        # n*1024
        self.fc8 = nn.Sequential(
            nn.Linear(1024, 75))
        # n*75

    def forward(self, x):
        x = self.c1(x)
        x = self.c2(x)
        x = self.c3(x)
        x = self.c4(x)
        x = self.c5(x)
        x = x.view(-1, 64 * 6 * 6)
        x = self.fc6(x)
        x = self.fc7(x)
        x = self.fc8(x)
        return x
